Knowledge bases for Amazon Bedrock を HashiCorp Terraform で作ってみる

HashiCorp Terraform の記事がなかったのでブログ化してみました。 Aurora は停止ができるため、不要な時間は停止しましょう。
#Amazon Bedrock Knowledge Bases
#AWS
#Amazon Bedrock
#生成AI
たかくに
2024.04.27
こんにちは！ AWS 事業本部コンサルティング部のたかくに（@takakuni_） です。
今回はタイトルの通り、 Knowledge bases for Amazon Bedrock を HashiCorp Terraform で作成してみようと思います。
 はじめにHashiCorp Terraform で、 AWS をデプロイする場合、多くのケースでは Terraform AWS Provider （リソース名が aws_ で始まるもの）が使われるかと思います。
2024/04/26 にリリースされた Terraform AWS Provider v5.47.0 では、 Knowledge bases for Amazon Bedrock のデータソースがデプロイできないため、 Terraform AWS Cloud Control Provider を組み合わせる必要があります。
実は直近の Terraform AWS Provider v5.46.0 まで、ナレッジベース本体もサポートしていなかったため、 v5.47.0 のアップデートが個人的には少し嬉しかったです。
そこで、今回は次のプロバイダーバージョンを利用して、ナレッジベースを作成していきます。
Terraform AWS Provider v5.47.0
ナレッジベース本体とナレッジベースで利用する IAM ロールを作成
Terraform AWS Cloud Control Provider v0.75.0
ナレッジベースのデータソースにあたる部分を作成 [1]
最終的なコードは以下に格納しています。
https://github.com/takakuni-classmethod/genai-blog/tree/main/knowledge_bases_support_terraform
 構成図今回の構成図は次のとおりです。各構成要素に対して考慮ポイントを押さえていきます。
 ベクトルデータベース 前提条件Amazon Aurora で提供されているデータベースエンジンのうち、 Knowledge bases for Amazon Bedrock で利用可能なものは現状 PostgreSQL のみです。
また、 Aurora PostgreSQL を ベクトルデータベースとして利用するにはいくつか前提条件があります。
バージョン
15.4 以上
14.9 以上
13.12 以上
12.16 以上
Data API が必要
Secrets Manager でパスワード管理
Knowledge bases for Amazon Bedrock で利用
Using Aurora PostgreSQL as a Knowledge Base for Amazon Bedrock - Amazon Aurora
 データベースのセットアップ加えて、ベクトルデータベースで利用するための前段階として、いくつかクエリを実行する必要があります。
Using Aurora PostgreSQL as a Knowledge Base for Amazon Bedrock - Amazon Aurora
今回、 SQL のクエリは Terraform の local-exec Provisioner を利用し Data API 経由でクエリを実行しました。
万が一クエリの実行が重複しても問題ないように、ドキュメントに記載のあるクエリを少しばかり改修してみました。（このあたりあまり明るくないので、もっとこうした方がいいとかあれば、優しく教えていただけると嬉しいです）
#!/bin/bash
# For more information, see the official documentation:
# https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/AuroraPostgreSQL.VectorDB.html

ROLE_NAME=$(echo $ROLE_NAME | sed 's/"//g')
PASSWORD=$(echo $PASSWORD | sed 's/"//g')
EMBEDDING_DIMENSION=$(echo $EMBEDDING_DIMENSION | sed 's/"//g')

# Install pgvector extension
aws rds-data --region $REGION execute-statement \
  --resource-arn $CLUSTER_ARN \
  --secret-arn $SECRET_ARN \
  --database $DATABASE_NAME \
  --sql "CREATE EXTENSION IF NOT EXISTS vector;"

# Check the version of the pg_vector extension
aws rds-data --region $REGION execute-statement \
  --resource-arn $CLUSTER_ARN \
  --secret-arn $SECRET_ARN \
  --database $DATABASE_NAME \
  --sql "SELECT extversion FROM pg_extension WHERE extname='vector';"

# Create a specific schema that Bedrock can use to query the data
aws rds-data --region $REGION execute-statement \
  --resource-arn $CLUSTER_ARN \
  --secret-arn $SECRET_ARN \
  --database $DATABASE_NAME \
  --sql "CREATE SCHEMA IF NOT EXISTS bedrock_integration;"

# Create a new role that Bedrock can use to query the database

aws rds-data --region $REGION execute-statement \
  --resource-arn $CLUSTER_ARN \
  --secret-arn $SECRET_ARN \
  --database $DATABASE_NAME \
  --sql "
    DO
    \$do\$
    BEGIN
      IF EXISTS (
          SELECT FROM pg_catalog.pg_roles 
          WHERE rolname = '$ROLE_NAME') THEN

          RAISE NOTICE 'Role $ROLE_NAME already exists. Skipping.';
      ELSE
          CREATE ROLE $ROLE_NAME LOGIN PASSWORD '$PASSWORD';
      END IF;
    END
    \$do\$;"

# Grant the bedrock_user permission to manage the bedrock_integration schema
aws rds-data --region $REGION execute-statement \
  --resource-arn $CLUSTER_ARN \
  --secret-arn $SECRET_ARN \
  --database $DATABASE_NAME \
  --sql "GRANT ALL ON SCHEMA bedrock_integration to $ROLE_NAME;"

# Login as the bedrock_user and create a table in the bedrock_integration schema
aws rds-data --region $REGION execute-statement \
  --resource-arn $CLUSTER_ARN \
  --secret-arn $SECRET_ARN \
  --database $DATABASE_NAME \
  --sql "CREATE TABLE IF NOT EXISTS bedrock_integration.bedrock_kb (id uuid PRIMARY KEY, embedding vector($EMBEDDING_DIMENSION), chunks text, metadata json);"

# Create an index with the cosine operator which the bedrock can use to query the data
aws rds-data --region $REGION execute-statement \
  --resource-arn $CLUSTER_ARN \
  --secret-arn $SECRET_ARN \
  --database $DATABASE_NAME \
  --sql "CREATE INDEX IF NOT EXISTS bedrock_kb_embedding_idx on bedrock_integration.bedrock_kb USING hnsw (embedding vector_cosine_ops);"
 Aurora Serverless v2 の理由続いてインスタンスタイプの選定です。今回は単純にコスト意識で Aurora Serverless v2 を選定しました。
Knowledge bases for Amazon Bedrock は Data API 経由で Embedding を行います。
オンデマンドインスタンスの場合、 T 系インスタンスでは、 Data API をサポートしていないため、最低でも R6g 系の利用、フル稼働で最低でも月額 211 ドル前後かかってしまいます。
Data API は T DB インスタンスクラスではサポートされていません。
RDS Data API の使用 - Amazon Aurora
個人利用であれば、 Aurora Serverless v2 の 0.5 ACU、月額 43.80 ドルからスタートするのがお手頃かと思います。
 Data API についてPostgreSQL であれば、通常 5432 ポートで接続すると思いますが、 AWS では Data API といったセキュアな HTTP エンドポイント経由で接続できる方法が設けられています。
Serverless v2 とプロビジョニングされた Amazon Aurora PostgreSQL で RDS Data API がサポートされました | DevelopersIO
今回の構成の場合、 Aurora へのアクセスは Data API 経由で完結するため、セキュリティグループにインバウンドルールがない状態でデータベースが完成します。
Data API はとても魅力的ですが、別途クエリに対しコストが発生するのと、実行できるクエリにもサポート可否があるため利用は計画的に行いましょう。
 Knowledge bases for Amazon Bedrock モデル エンべディングモデルの次元数を取得したい再掲になりますが、ベクトルデータベースのセットアップで以下のクエリのようにスキーマを定義してあげる必要があります。
# Login as the bedrock_user and create a table in the bedrock_integration schema
aws rds-data --region $REGION execute-statement \
  --resource-arn $CLUSTER_ARN \
  --secret-arn $SECRET_ARN \
  --database $DATABASE_NAME \
  --sql "CREATE TABLE IF NOT EXISTS bedrock_integration.bedrock_kb (id uuid PRIMARY KEY, embedding vector($EMBEDDING_DIMENSION), chunks text, metadata json);"
実はこのクエリで指定するモデルの次元数(EMBEDDING_DIMENSION)が、エンべディングモデルの種類によって異なります。
Titan Embeddings G1 - Text：1536 次元
Titan Multimodal Embeddings G1：1,024 (デフォルト), 384, または 256 次元
Cohere Embed English：1024 次元
Cohere Embed Multilingual：1024 次元
おそらく普段から Terraform を使っている方なら、 data ブロックを使って以下のように定義するかと思います。
data "aws_bedrock_foundation_model" "embedding" {
  model_id = var.knowledge_bases.embeddings_model_id
}
しかし、 data ブロックにはモデルの次元数が含まれておらず、以下の値が取れてきます。
# data.aws_bedrock_foundation_model.embedding:
data "aws_bedrock_foundation_model" "embedding" {
    customizations_supported     = []
    id                           = "cohere.embed-multilingual-v3"
    inference_types_supported    = [
        "ON_DEMAND",
    ]
    input_modalities             = [
        "TEXT",
    ]
    model_arn                    = "arn:aws:bedrock:us-east-1::foundation-model/cohere.embed-multilingual-v3"
    model_id                     = "cohere.embed-multilingual-v3"
    model_name                   = "Embed Multilingual"
    output_modalities            = [
        "EMBEDDING",
    ]
    provider_name                = "Cohere"
    response_streaming_supported = false
}
AWS CLI でも引っ張れないようなので、マネコンから目で確認して静的な値を埋め込む必要がありました。今後 API 経由で引っ張れることに期待ですね。
 ナレッジベースナレッジベースは、 Terraform AWS Provider を使います。以下のように定義してあげることで実装できました。
########################################################
# Knowledge Bases
########################################################
resource "aws_bedrockagent_knowledge_base" "this" {
  name     = "${local.prefix}-kb"
  role_arn = aws_iam_role.knowledge_bases.arn

  knowledge_base_configuration {
    type = "VECTOR"
    vector_knowledge_base_configuration {
      embedding_model_arn = data.aws_bedrock_foundation_model.embedding.model_arn
    }
  }

  storage_configuration {
    type = "RDS"
    rds_configuration {
      credentials_secret_arn = module.vector_db.secrets.arn
      database_name          = module.vector_db.cluster.database_name
      resource_arn           = module.vector_db.cluster.arn
      table_name             = "bedrock_integration.bedrock_kb"
      field_mapping {
        primary_key_field = "id"
        vector_field      = "embedding"
        text_field        = "chunks"
        metadata_field    = "metadata"
      }
    }
  }

  depends_on = [aws_iam_role_policy_attachment.knowledge_bases]
}
余談ですが Terraform AWS Cloud Control Provider の時は、ブロック {} が、オブジェクト = {} で定義されており、型が若干違ってて面白かったです。
########################################################
# Knowledge Base
########################################################
resource "awscc_bedrock_knowledge_base" "this" {
  name     = "${local.prefix}-kb"
  role_arn = aws_iam_role.knowledge_bases.arn

  knowledge_base_configuration = {
    type = "VECTOR"
    vector_knowledge_base_configuration = {
      embedding_model_arn = data.aws_bedrock_foundation_model.embedding.model_arn
    }
  }

  storage_configuration = {
    type = "RDS"
    rds_configuration = {
      credentials_secret_arn = module.vector_db.secrets.arn
      database_name          = module.vector_db.cluster.database_name
      resource_arn           = module.vector_db.cluster.arn
      table_name             = "bedrock_integration.bedrock_kb"
      field_mapping = {
        primary_key_field = "id"
        vector_field      = "embedding"
        text_field        = "chunks"
        metadata_field    = "metadata"
      }
    }
  }

  depends_on = [aws_iam_role_policy_attachment.knowledge_bases]
}
 データソースデータソースは Terraform AWS Cloud Control Provider を利用して定義します。 aws_bedrockagent_knowledge_base の Attribute Reference に id が書いてなかったのですが、値が取れてよかったです。
########################################################
# Knowledge Base Data Source
########################################################
resource "awscc_bedrock_data_source" "this" {
  name              = local.prefix
  knowledge_base_id = aws_bedrockagent_knowledge_base.this.id
  data_source_configuration = {
    type = "S3"
    s3_configuration = {
      bucket_arn = module.datasource.bucket.arn
    }
  }
}
 動作確認Terraform のコード解説は以上です。詳しくみたい方は GitHub をみていただけると嬉しいです。
それでは、きちんと同期できるか確認します。まずはその辺りに転がっていた就業規則を持ってきます。（嘘です。生成 AI に用意してもらいました）
# 就業規則

## 第 1 章 総則

### 第 1 条 目的

本規則は、株式会社 ○○（以下「会社」という）の従業員の就業に関する事項を定めるものである。

### 第 2 条 適用範囲

本規則は、会社のすべての従業員に適用する。

## 第 2 章 採用・異動

### 第 3 条 採用

会社は、必要な人材を公平に選考し、採用する。

### 第 4 条 異動

会社は、業務上必要がある場合、従業員の異動を命ずることがある。

## 第 3 章 勤務

### 第 5 条 勤務時間

従業員の勤務時間は、原則として次のとおりとする。
（1）始業時刻 09 時 00 分
（2）終業時刻 17 時 30 分

### 第 6 条 休憩時間

休憩時間は、12 時 00 分から 13 時 00 分までの 60 分とする。

### 第 7 条 休日

休日は、土曜日、日曜日、国民の祝日および会社の指定する日とする。

## 第 4 章 服務

### 第 8 条 服務規律

従業員は、会社の諸規則を遵守し、上司の指示に従って、誠実に職務を遂行しなければならない。

### 第 9 条 秘密保持

従業員は、在職中または退職後においても、会社の秘密情報を漏らしてはならない。

## 第 5 章 賃金

### 第 10 条 賃金

賃金は、別に定める賃金規定による。

## 第 6 章 安全衛生

### 第 11 条 安全衛生

会社は、従業員の安全と健康を確保するため、必要な措置を講ずる。

## 第 7 章 懲戒

### 第 12 条 懲戒

従業員が本規則に違反した場合、会社は懲戒処分を行うことがある。

## 附則

本規則は、○○ 年 ○ 月 ○ 日から施行する。
該当のデータソースバケットに先ほどのマークダウンをアップロードします。
Bedrock コンソールから同期ボタンで同期を取ってみます。
1 つのファイルなので、ほとんど時間かからず同期が取れましたね。回答もしっかり返ってきています。
クエリエディタから Aurora PostgreSQL のデータもみてみます。
SELECT * FROM bedrock_integration.bedrock_kb LIMIT 5;
embedding 列に 1024 の数字ベクトルが続いてました。ちゃんとエンべディングできていますね。
 まとめ以上、簡単ではありますが「Knowledge bases for Amazon Bedrock を HashiCorp Terraform で作ってみる」でした。
現状、 2 つのプロバイダーを使う必要がありますが、今後のアップデートに期待ですね。このブログが皆様の参考になれば幸いです。
AWS 事業本部コンサルティング部のたかくに（@takakuni_）でした！

脚注
Terraform AWS Cloud Control Provider のみでも Knowledge bases for Amazon Bedrock の一連のリソースは作成可能です。いち早く新しい機能を Terraform でデプロイしたい場合は、 awscc を追ってみるといいと思います。 ↩︎